#!/bin/bash

# ES Archive 数据同步脚本
# 使用 Docker 容器进行 Elasticsearch 数据备份和恢复

set -e

# 默认配置
DEFAULT_ES_HOST="http://elasticsearch:9200"
DEFAULT_BACKUP_DIR="$(pwd)/backups"
DEFAULT_METADATA_DIR="$(pwd)/metadata"
DEFAULT_LOGS_DIR="$(pwd)/logs"
DEFAULT_CONFIG_DIR="$(pwd)/config"
DEFAULT_IMAGE="es-archive:minimal"
DEFAULT_NETWORK=""

# 颜色输出
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
NC='\033[0m' # No Color

# 打印帮助信息
print_help() {
    echo "ES Archive 数据同步工具"
    echo ""
    echo "用法:"
    echo "  $0 backup <index> [options]     # 备份数据"
    echo "  $0 restore <backup_id> <index> [options]  # 恢复数据"
    echo "  $0 list                         # 列出备份"
    echo "  $0 validate <backup_id>         # 验证备份"
    echo ""
    echo "环境变量:"
    echo "  ES_HOST       Elasticsearch 地址 (默认: $DEFAULT_ES_HOST)"
    echo "  ES_USERNAME   Elasticsearch 用户名"
    echo "  ES_PASSWORD   Elasticsearch 密码"
    echo "  BACKUP_DIR    备份目录 (默认: $DEFAULT_BACKUP_DIR)"
    echo "  DOCKER_IMAGE  Docker 镜像 (默认: $DEFAULT_IMAGE)"
    echo "  DOCKER_NETWORK Docker 网络"
    echo ""
    echo "示例:"
    echo "  # 备份所有日志索引"
    echo "  $0 backup 'logs-*'"
    echo ""
    echo "  # 使用配置文件备份"
    echo "  $0 backup 'app-*' --config config/default.yaml"
    echo ""
    echo "  # 备份指定时间范围的数据"
    echo "  ES_HOST=http://es:9200 $0 backup app-logs --time-range '2025-07-01' '2025-07-31'"
    echo ""
    echo "  # 恢复数据"
    echo "  $0 restore backup_20250731_001 restored-logs"
    echo ""
    echo "  # 列出所有备份"
    echo "  $0 list"
}

# 日志函数
log_info() {
    echo -e "${BLUE}[INFO]${NC} $1"
}

log_success() {
    echo -e "${GREEN}[SUCCESS]${NC} $1"
}

log_warning() {
    echo -e "${YELLOW}[WARNING]${NC} $1"
}

log_error() {
    echo -e "${RED}[ERROR]${NC} $1"
}

# 检查必要的目录
ensure_directories() {
    local backup_dir="${BACKUP_DIR:-$DEFAULT_BACKUP_DIR}"
    local metadata_dir="${METADATA_DIR:-$DEFAULT_METADATA_DIR}"
    local logs_dir="${LOGS_DIR:-$DEFAULT_LOGS_DIR}"
    
    mkdir -p "$backup_dir" "$metadata_dir" "$logs_dir"
    log_info "创建目录: $backup_dir, $metadata_dir, $logs_dir"
}

# 构建 Docker 运行命令的基础部分
build_docker_base() {
    local es_host="${ES_HOST:-$DEFAULT_ES_HOST}"
    local backup_dir="${BACKUP_DIR:-$DEFAULT_BACKUP_DIR}"
    local metadata_dir="${METADATA_DIR:-$DEFAULT_METADATA_DIR}"
    local logs_dir="${LOGS_DIR:-$DEFAULT_LOGS_DIR}"
    local config_dir="${CONFIG_DIR:-$DEFAULT_CONFIG_DIR}"
    local image="${DOCKER_IMAGE:-$DEFAULT_IMAGE}"
    local network="${DOCKER_NETWORK:-$DEFAULT_NETWORK}"
    
    local docker_cmd="docker run --rm"
    
    # 添加卷挂载
    docker_cmd="$docker_cmd -v $backup_dir:/data/backups"
    docker_cmd="$docker_cmd -v $metadata_dir:/data/metadata"
    docker_cmd="$docker_cmd -v $logs_dir:/data/logs"
    
    # 如果配置目录存在，则挂载
    if [ -d "$config_dir" ]; then
        docker_cmd="$docker_cmd -v $config_dir:/app/config"
    fi
    
    # 添加环境变量
    docker_cmd="$docker_cmd -e ES_HOST=$es_host"
    
    if [ -n "$ES_USERNAME" ]; then
        docker_cmd="$docker_cmd -e ES_USERNAME=$ES_USERNAME"
    fi
    
    if [ -n "$ES_PASSWORD" ]; then
        docker_cmd="$docker_cmd -e ES_PASSWORD=$ES_PASSWORD"
    fi
    
    if [ -n "$ES_API_KEY" ]; then
        docker_cmd="$docker_cmd -e ES_API_KEY=$ES_API_KEY"
    fi
    
    # 添加网络
    if [ -n "$network" ]; then
        docker_cmd="$docker_cmd --network $network"
    fi
    
    # 添加镜像
    docker_cmd="$docker_cmd $image"
    
    echo "$docker_cmd"
}

# 备份数据
backup_data() {
    local index="$1"
    shift
    
    if [ -z "$index" ]; then
        log_error "请指定要备份的索引名称"
        exit 1
    fi
    
    ensure_directories
    
    local docker_cmd=$(build_docker_base)
    local backup_cmd="es-archive backup '$index' /data/backups"
    
    # 解析额外参数
    while [[ $# -gt 0 ]]; do
        case $1 in
            --time-range)
                if [ -n "$2" ] && [ -n "$3" ]; then
                    backup_cmd="$backup_cmd --strategy time_range --start-time $2 --end-time $3"
                    shift 3
                else
                    log_error "--time-range 需要两个参数: 开始时间和结束时间"
                    exit 1
                fi
                ;;
            --batch-size)
                backup_cmd="$backup_cmd --batch-size $2"
                shift 2
                ;;
            --compression)
                backup_cmd="$backup_cmd --compression $2"
                shift 2
                ;;
            --no-validate)
                backup_cmd="$backup_cmd --no-validate"
                shift
                ;;
            *)
                backup_cmd="$backup_cmd $1"
                shift
                ;;
        esac
    done
    
    log_info "开始备份索引: $index"
    log_info "执行命令: $docker_cmd $backup_cmd"
    
    if eval "$docker_cmd $backup_cmd"; then
        log_success "备份完成"
    else
        log_error "备份失败"
        exit 1
    fi
}

# 恢复数据
restore_data() {
    local backup_id="$1"
    local target_index="$2"
    shift 2
    
    if [ -z "$backup_id" ] || [ -z "$target_index" ]; then
        log_error "请指定备份ID和目标索引名称"
        exit 1
    fi
    
    ensure_directories
    
    local docker_cmd=$(build_docker_base)
    local restore_cmd="es-archive restore '$backup_id' '$target_index'"
    
    # 解析额外参数
    while [[ $# -gt 0 ]]; do
        case $1 in
            --overwrite)
                restore_cmd="$restore_cmd --overwrite"
                shift
                ;;
            --no-validate)
                restore_cmd="$restore_cmd --no-validate"
                shift
                ;;
            --batch-size)
                restore_cmd="$restore_cmd --batch-size $2"
                shift 2
                ;;
            *)
                restore_cmd="$restore_cmd $1"
                shift
                ;;
        esac
    done
    
    log_info "开始恢复数据: $backup_id -> $target_index"
    log_info "执行命令: $docker_cmd $restore_cmd"
    
    if eval "$docker_cmd $restore_cmd"; then
        log_success "恢复完成"
    else
        log_error "恢复失败"
        exit 1
    fi
}

# 列出备份
list_backups() {
    ensure_directories
    
    local docker_cmd=$(build_docker_base)
    local list_cmd="es-archive list-backups"
    
    log_info "列出所有备份"
    eval "$docker_cmd $list_cmd"
}

# 验证备份
validate_backup() {
    local backup_id="$1"
    
    if [ -z "$backup_id" ]; then
        log_error "请指定要验证的备份ID"
        exit 1
    fi
    
    ensure_directories
    
    local docker_cmd=$(build_docker_base)
    local validate_cmd="es-archive validate '$backup_id' --detailed"
    
    log_info "验证备份: $backup_id"
    eval "$docker_cmd $validate_cmd"
}

# 主函数
main() {
    case "${1:-}" in
        backup)
            shift
            backup_data "$@"
            ;;
        restore)
            shift
            restore_data "$@"
            ;;
        list)
            list_backups
            ;;
        validate)
            shift
            validate_backup "$@"
            ;;
        help|--help|-h)
            print_help
            ;;
        *)
            log_error "未知命令: ${1:-}"
            echo ""
            print_help
            exit 1
            ;;
    esac
}

# 执行主函数
main "$@"
